#include "yolo5_pre.h"
#include <iostream>
#include <fstream>
#include <filesystem>
#include <algorithm>
#include <cstring>
#include <opencv2/core.hpp>

namespace fs = std::filesystem;

Yolo5Pre::Yolo5Pre() {
    m_rga2Wrapper = std::make_unique<RGA2Wrapper>();
    if (!m_rga2Wrapper->initialize()) {
        std::cerr << "警告: RGA2初始化失败，将使用OpenCV进行图像处理" << std::endl;
    }
}

Yolo5Pre::~Yolo5Pre() = default;

bool Yolo5Pre::processDataset(const std::string& sourceDir, const std::string& destDir, double rotationAngle) {
    std::cout << "源目录: " << sourceDir << std::endl;
    std::cout << "目标目录: " << destDir << std::endl;
    if (rotationAngle != 0.0) {
        std::cout << "旋转角度: " << rotationAngle << " 度" << std::endl;
    }
    
    // 创建输出目录
    if (!createOutputDirectory(destDir)) {
        return false;
    }
    
    // 获取文件列表
    std::vector<std::string> rawFiles = getFileList(sourceDir, ".raw");
    std::vector<std::string> pngFiles = getFileList(sourceDir, ".png");
    
    std::cout << "发现 " << (rawFiles.size() + pngFiles.size()) << " 个文件:" << std::endl;
    std::cout << "  .raw文件: " << rawFiles.size() << " 个" << std::endl;
    std::cout << "  .png文件: " << pngFiles.size() << " 个" << std::endl;
    
    // 处理.raw文件
    if (!rawFiles.empty()) {
        std::cout << "\n开始转换NV12格式的.raw文件到640x640..." << std::endl;
        int convertedCount = 0;
        for (const auto& rawFile : rawFiles) {
            std::string rawPath = (fs::path(sourceDir) / rawFile).string();
            if (convertNV12To640x640PNG(rawPath, destDir, rotationAngle)) {
                convertedCount++;
            }
        }
        std::cout << "转换完成: " << convertedCount << "/" << rawFiles.size() << " 个文件" << std::endl;
    }
    
    // 处理PNG文件
    if (!pngFiles.empty()) {
        std::cout << "\n开始填充PNG图片到640x640..." << std::endl;
        int paddedCount = 0;
        for (const auto& pngFile : pngFiles) {
            std::string pngPath = (fs::path(sourceDir) / pngFile).string();
            if (padImageTo640x640(pngPath, destDir, rotationAngle)) {
                paddedCount++;
            }
        }
        std::cout << "填充完成: " << paddedCount << "/" << pngFiles.size() << " 个文件" << std::endl;
    }
    
    if (rawFiles.empty() && pngFiles.empty()) {
        std::cout << "  未发现.raw或.png文件" << std::endl;
    }
    
    return true;
}

bool Yolo5Pre::convertNV12To640x640PNG(const std::string& rawFilePath, const std::string& outputDir, double rotationAngle) {
    try {
        using namespace Yolo5PreConstants;
        
        // 读取.raw文件
        std::ifstream file(rawFilePath, std::ios::binary);
        if (!file.is_open()) {
            std::cerr << "  无法打开文件: " << rawFilePath << std::endl;
            return false;
        }
        
        // 获取文件大小
        file.seekg(0, std::ios::end);
        size_t fileSize = file.tellg();
        file.seekg(0, std::ios::beg);
        
        // 验证文件大小
        if (fileSize < NV12_TOTAL_SIZE) {
            std::cout << "  警告: " << fs::path(rawFilePath).filename().string() 
                      << " 文件大小不足，期望 " << NV12_TOTAL_SIZE << " 字节，实际 " << fileSize << " 字节，跳过处理" << std::endl;
            return false;
        }
        
        // 读取数据
        std::vector<uint8_t> rawData(fileSize);
        file.read(reinterpret_cast<char*>(rawData.data()), fileSize);
        file.close();
        
        // 从NV12数据创建OpenCV Mat (480x480 BGR格式)
        cv::Mat bgrMat = createMatFromNV12(rawData, INPUT_WIDTH, INPUT_HEIGHT);
        if (bgrMat.empty()) {
            std::cerr << "  无法从NV12数据创建图像" << std::endl;
            return false;
        }
        
        // 验证转换后的图像尺寸
        if (bgrMat.rows != INPUT_HEIGHT || bgrMat.cols != INPUT_WIDTH) {
            std::cerr << "  图像尺寸错误: 期望 " << INPUT_WIDTH << "x" << INPUT_HEIGHT 
                      << "，实际 " << bgrMat.cols << "x" << bgrMat.rows << std::endl;
            return false;
        }
        
        // 创建640x640的输出图像
        cv::Mat outputMat;
        if (rotationAngle != 0.0) {
            // 先旋转，再填充
            if (!m_rga2Wrapper->rotateAndPad(bgrMat, outputMat, rotationAngle, OUTPUT_SIZE, OUTPUT_SIZE)) {
                // 如果RGA2失败，回退到OpenCV
                cv::Mat rotated;
                cv::Point2f center(INPUT_WIDTH / 2.0f, INPUT_HEIGHT / 2.0f);
                cv::Mat rotationMatrix = cv::getRotationMatrix2D(center, rotationAngle, 1.0);
                cv::warpAffine(bgrMat, rotated, rotationMatrix, cv::Size(INPUT_WIDTH, INPUT_HEIGHT));
                
                outputMat = cv::Mat::zeros(OUTPUT_SIZE, OUTPUT_SIZE, rotated.type());
                rotated.copyTo(outputMat(cv::Rect(0, 0, INPUT_WIDTH, INPUT_HEIGHT)));
            }
        } else {
            // 直接填充
            if (!m_rga2Wrapper->padImage(bgrMat, outputMat, OUTPUT_SIZE, OUTPUT_SIZE)) {
                // 如果RGA2失败，回退到OpenCV
                outputMat = cv::Mat::zeros(OUTPUT_SIZE, OUTPUT_SIZE, bgrMat.type());
                bgrMat.copyTo(outputMat(cv::Rect(0, 0, INPUT_WIDTH, INPUT_HEIGHT)));
            }
        }
        
        // 生成输出文件名
        std::string outputFilename = generateOutputFilename(rawFilePath, outputDir, rotationAngle, "_640x640.png");
        std::string outputPath = (fs::path(outputDir) / outputFilename).string();
        
        // 保存PNG文件
        if (cv::imwrite(outputPath, outputMat)) {
            std::cout << "  转换完成: " << fs::path(rawFilePath).filename().string() 
                      << " -> " << outputFilename << " (" << OUTPUT_SIZE << "x" << OUTPUT_SIZE << ")" << std::endl;
            return true;
        } else {
            std::cerr << "  保存失败: " << outputPath << std::endl;
            return false;
        }
        
    } catch (const std::exception& e) {
        std::cerr << "  转换失败: " << rawFilePath << " - " << e.what() << std::endl;
        return false;
    }
}

bool Yolo5Pre::padImageTo640x640(const std::string& imagePath, const std::string& outputDir, double rotationAngle) {
    try {
        using namespace Yolo5PreConstants;
        
        // 读取图片
        cv::Mat img = cv::imread(imagePath);
        if (img.empty()) {
            std::cerr << "  读取失败: " << imagePath << std::endl;
            return false;
        }
        
        // 检查图片尺寸
        int h = img.rows, w = img.cols;
        if (h != INPUT_HEIGHT || w != INPUT_WIDTH) {
            std::cout << "  跳过: " << fs::path(imagePath).filename().string() 
                      << " 尺寸为 " << w << "x" << h << "，不是" << INPUT_WIDTH << "x" << INPUT_HEIGHT << std::endl;
            return false;
        }
        
        cv::Mat outputMat;
        if (rotationAngle != 0.0) {
            // 先旋转，再填充
            if (!m_rga2Wrapper->rotateAndPad(img, outputMat, rotationAngle, OUTPUT_SIZE, OUTPUT_SIZE)) {
                // 如果RGA2失败，回退到OpenCV
                cv::Mat rotated;
                cv::Point2f center(w / 2.0f, h / 2.0f);
                cv::Mat rotationMatrix = cv::getRotationMatrix2D(center, rotationAngle, 1.0);
                cv::warpAffine(img, rotated, rotationMatrix, cv::Size(w, h));
                
                outputMat = cv::Mat::zeros(OUTPUT_SIZE, OUTPUT_SIZE, rotated.type());
                rotated.copyTo(outputMat(cv::Rect(0, 0, w, h)));
            }
            std::cout << "  已旋转 " << rotationAngle << " 度" << std::endl;
        } else {
            // 直接填充
            if (!m_rga2Wrapper->padImage(img, outputMat, OUTPUT_SIZE, OUTPUT_SIZE)) {
                // 如果RGA2失败，回退到OpenCV
                outputMat = cv::Mat::zeros(OUTPUT_SIZE, OUTPUT_SIZE, img.type());
                img.copyTo(outputMat(cv::Rect(0, 0, w, h)));
            }
        }
        
        // 生成输出文件名
        std::string outputFilename = generateOutputFilename(imagePath, outputDir, rotationAngle, "_640x640.png");
        std::string outputPath = (fs::path(outputDir) / outputFilename).string();
        
        // 保存填充后的图片
        if (cv::imwrite(outputPath, outputMat)) {
            std::cout << "  填充完成: " << fs::path(imagePath).filename().string() 
                      << " -> " << outputFilename << std::endl;
            return true;
        } else {
            std::cerr << "  保存失败: " << outputPath << std::endl;
            return false;
        }
        
    } catch (const std::exception& e) {
        std::cerr << "  填充失败: " << imagePath << " - " << e.what() << std::endl;
        return false;
    }
}

bool Yolo5Pre::validateDirectory(const std::string& directoryPath) {
    if (!fs::exists(directoryPath)) {
        std::cerr << "错误: 目录 '" << directoryPath << "' 不存在" << std::endl;
        return false;
    }
    
    if (!fs::is_directory(directoryPath)) {
        std::cerr << "错误: '" << directoryPath << "' 不是一个目录" << std::endl;
        return false;
    }
    
    return true;
}

cv::Mat Yolo5Pre::createMatFromNV12(const std::vector<uint8_t>& rawData, int width, int height) {
    using namespace Yolo5PreConstants;
    
    // 固定输入尺寸验证
    if (width != INPUT_WIDTH || height != INPUT_HEIGHT) {
        std::cerr << "   错误: 不支持的尺寸 " << width << "x" << height << "，仅支持" << INPUT_WIDTH << "x" << INPUT_HEIGHT << std::endl;
        return cv::Mat();
    }
    
    if (rawData.size() < NV12_TOTAL_SIZE) {
        std::cerr << "    NV12数据大小不足: 期望 " << NV12_TOTAL_SIZE << " 字节，实际 " << rawData.size() << " 字节" << std::endl;
        return cv::Mat();
    }
    
    std::cout << "   处理NV12数据: " << width << "x" << height << ", Y分量: " << NV12_Y_SIZE << " 字节, UV分量: " << NV12_UV_SIZE << " 字节" << std::endl;
    
    try {
        // 创建完整的YUV矩阵 (height * 3/2 × width)
        // NV12格式：Y分量 + UV分量交错存储
        cv::Mat yuvMat(height * 3 / 2, width, CV_8UC1);
        
        // 复制Y分量 (前480行)
        std::memcpy(yuvMat.data, rawData.data(), NV12_Y_SIZE);
        
        // 复制UV分量 (后240行)
        std::memcpy(yuvMat.data + NV12_Y_SIZE, rawData.data() + NV12_Y_SIZE, NV12_UV_SIZE);
        
        std::cout << "   YUV矩阵创建成功: " << yuvMat.rows << "x" << yuvMat.cols << ", 类型: " << yuvMat.type() << std::endl;
        
        // 使用OpenCV的NV12到BGR转换
        cv::Mat bgrMat;
        cv::cvtColor(yuvMat, bgrMat, cv::COLOR_YUV2BGR_NV12);
        
        if (bgrMat.empty()) {
            std::cerr << "    OpenCV颜色空间转换失败" << std::endl;
            return cv::Mat();
        }
        
        std::cout << "   BGR转换成功: " << bgrMat.rows << "x" << bgrMat.cols << ", 类型: " << bgrMat.type() << std::endl;
        
        return bgrMat;
        
    } catch (const cv::Exception& e) {
        std::cerr << "    OpenCV异常: " << e.what() << std::endl;
        return cv::Mat();
    } catch (const std::exception& e) {
        std::cerr << "    异常: " << e.what() << std::endl;
        return cv::Mat();
    }
}

std::string Yolo5Pre::generateOutputFilename(const std::string& inputPath, const std::string& outputDir, 
                                           double rotationAngle, const std::string& suffix) {
    (void)outputDir; // 避免未使用参数警告
    std::string baseName = fs::path(inputPath).stem().string();
    std::string outputFilename;
    
    if (rotationAngle != 0.0) {
        outputFilename = baseName + "_rot" + std::to_string(static_cast<int>(rotationAngle)) + suffix;
    } else {
        outputFilename = baseName + suffix;
    }
    
    return outputFilename;
}

bool Yolo5Pre::createOutputDirectory(const std::string& outputDir) {
    try {
        fs::create_directories(outputDir);
        std::cout << "输出目录已创建: " << outputDir << std::endl;
        return true;
    } catch (const std::exception& e) {
        std::cerr << "错误: 无法创建输出目录 '" << outputDir << "': " << e.what() << std::endl;
        return false;
    }
}

std::vector<std::string> Yolo5Pre::getFileList(const std::string& directory, const std::string& extension) {
    std::vector<std::string> files;
    
    try {
        for (const auto& entry : fs::directory_iterator(directory)) {
            if (entry.is_regular_file() && entry.path().extension() == extension) {
                files.push_back(entry.path().filename().string());
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "错误: 无法读取目录内容: " << e.what() << std::endl;
    }
    
    return files;
}
